import os
import json
import requests
from typing import List, Dict, Any, Union, Optional
import base64
from io import BytesIO
from PIL import Image

class VLLMModel:
    """与Ollama交互的VLLM模型封装类"""
    
    def __init__(self, base_url: str = "http://localhost:11434", model_name: str = "qwen2.5-vl"):
        """
        初始化VLLM模型
        
        Args:
            base_url: Ollama API的基础URL
            model_name: 模型名称
        """
        self.base_url = base_url
        self.model_name = model_name
    
    def get_available_models(self) -> List[str]:
        """
        获取Ollama中可用的模型列表
        
        Returns:
            模型名称列表
        """
        try:
            response = requests.get(f"{self.base_url}/api/tags")
            if response.status_code == 200:
                models_data = response.json()
                return [model["name"] for model in models_data.get("models", [])]
            else:
                print(f"获取模型列表失败: HTTP {response.status_code}")
                return []
        except Exception as e:
            print(f"获取模型列表时发生错误: {str(e)}")
            return []
    
    def image_to_base64(self, image_path: str) -> Optional[str]:
        """
        将图片转换为base64编码，并缩放到256*256分辨率
        
        Args:
            image_path: 图片路径
            
        Returns:
            base64编码的图片字符串
        """
        try:
            with Image.open(image_path) as img:
                # 确保图片是RGB模式
                if img.mode != "RGB":
                    img = img.convert("RGB")
                
                # 缩放图片到256*256分辨率
                img = img.resize((256, 256), Image.LANCZOS)
                    
                buffer = BytesIO()
                img.save(buffer, format="JPEG", quality=85)  # 使用较高的质量设置
                return base64.b64encode(buffer.getvalue()).decode("utf-8")
        except Exception as e:
            print(f"图片转base64失败: {str(e)}")
            return None
    
    def detect_defect(self, image_path: str, product_name: str, knowledge_results: List[Dict]) -> Dict:
        """
        缺陷检测模式
        
        Args:
            image_path: 图片路径
            product_name: 产品名称
            knowledge_results: 知识库检索结果
            
        Returns:
            检测结果
        """
        # 转换图片为base64
        image_b64 = self.image_to_base64(image_path)
        if not image_b64:
            return {"success": False, "error": "无法处理图片"}
        
        # 构建知识库检索结果的描述
        retrieval_context = self._format_retrieval_results(knowledge_results)
        
        # 获取缺陷类型描述
        defect_descriptions = ""
        if hasattr(knowledge_results, "get_defect_features_description"):
            defect_descriptions = knowledge_results.get_defect_features_description(product_name)
        else:
            # 尝试从检索结果中提取缺陷类型
            defect_types = set()
            for result in knowledge_results:
                label = result.get("label", "").strip()
                if label and label != "unknown":
                    defect_types.add(label)
            
            if defect_types:
                defect_descriptions = f"{product_name}产品可能的缺陷类型包括: {', '.join(defect_types)}。"
                defect_descriptions += "请特别注意区分正常产品(good)和有缺陷产品的区别。"
        
        # 为NEU-DET数据集添加特殊处理
        special_instructions = ""
        if product_name.lower() == "neu-det":
            special_instructions = """
对于钢材表面缺陷检测，请特别注意以下几点：

1. 龟裂(crazing)通常表现为网状细小裂纹，在光照下会形成特殊的反光图案，与划痕(scratches)的线状特征有明显区别。

2. 夹杂物(inclusion)通常是点状或不规则形状的异物，与点蚀(pitted_surface)的凹坑特征不同，夹杂物可能呈现为凸起或嵌入物。

3. 斑块(patches)表现为颜色变化区域，通常没有明显的深度变化，而点蚀和夹杂物通常伴随表面形态的变化。

4. 点蚀表面(pitted_surface)的特征是表面有凹坑或孔洞，通常在光照下会有明显的阴影效果。

5. 划痕(scratches)是线状损伤，通常有明显的方向性，与龟裂的网状特征有本质区别。

请仔细分析图像中的视觉特征，确保准确区分这些不同类型的表面缺陷。"""
        
        # 构建提示词
        prompt = f"""你是一个产品表面缺陷检测专家，请仔细分析图片中的{product_name}产品，并检测是否存在表面缺陷。

{defect_descriptions}

{special_instructions}

以下是从相似图片知识库中检索到的结果供你参考:
{retrieval_context}

请特别注意:
1. 仔细区分正常产品(good)和有缺陷产品的区别
2. 只有当产品表面完全无缺陷时，才应判断为"good"
3. 如有任何可疑特征，应详细分析是否为缺陷

请分析图片并输出标准JSON格式结果，包括:
1. has_defect: 布尔值，表示是否存在缺陷
2. defect_type: 如存在缺陷，提供缺陷类型；如无缺陷，则为"good"
3. confidence: 置信度，0-1的浮点数
4. analysis: 简短分析说明（不超过100字）

只需要输出JSON格式结果，不要有其他回复。"""

        # 构建API请求
        payload = {
            "model": self.model_name,
            "prompt": prompt,
            "stream": False,
            "images": [image_b64],
            "format": "json"
        }
        
        try:
            response = requests.post(f"{self.base_url}/api/generate", json=payload)
            if response.status_code == 200:
                result = response.json()
                # 尝试解析JSON响应
                try:
                    return {
                        "success": True,
                        "response": json.loads(result["response"])
                    }
                except json.JSONDecodeError:
                    # 如果无法解析为JSON，返回原始文本
                    return {
                        "success": True,
                        "response": result["response"]
                    }
            else:
                return {"success": False, "error": f"API调用失败: HTTP {response.status_code}"}
        except Exception as e:
            return {"success": False, "error": f"调用模型时发生错误: {str(e)}"}
    
    def chat(self, image_path: str, user_query: str, product_name: str, knowledge_results: List[Dict]) -> Dict:
        """
        对话模式
        
        Args:
            image_path: 图片路径
            user_query: 用户问题
            product_name: 产品名称
            knowledge_results: 知识库检索结果
            
        Returns:
            对话响应
        """
        # 转换图片为base64
        image_b64 = self.image_to_base64(image_path)
        if not image_b64:
            return {"success": False, "error": "无法处理图片"}
        
        # 构建知识库检索结果的描述
        retrieval_context = self._format_retrieval_results(knowledge_results)
        
        # 获取缺陷类型描述
        defect_descriptions = ""
        if hasattr(knowledge_results, "get_defect_features_description"):
            defect_descriptions = knowledge_results.get_defect_features_description(product_name)
        else:
            # 尝试从检索结果中提取缺陷类型
            defect_types = set()
            for result in knowledge_results:
                label = result.get("label", "").strip()
                if label and label != "unknown":
                    defect_types.add(label)
            
            if defect_types:
                defect_descriptions = f"{product_name}产品可能的缺陷类型包括: {', '.join(defect_types)}。"
        
        # 为NEU-DET数据集添加特殊处理
        special_instructions = ""
        if product_name.lower() == "neu-det":
            special_instructions = """
对于钢材表面缺陷分析，请注意以下关键区别：

- 龟裂(crazing): 网状细小裂纹，形成不规则网格状或蜘蛛网状纹路
- 夹杂物(inclusion): 点状或不规则形状的异物，可能是凸起或嵌入物
- 斑块(patches): 表面颜色变化区域，通常没有明显的深度变化
- 点蚀表面(pitted_surface): 表面有凹坑或孔洞，光照下有明显阴影
- 划痕(scratches): 线状损伤，有明显方向性

请根据这些特征进行分析，并在回答中明确指出观察到的特征如何对应到特定缺陷类型。"""
        
        # 构建提示词
        prompt = f"""你是一个产品表面缺陷检测专家，请根据用户的问题，分析图片中的{product_name}产品。

{defect_descriptions}

{special_instructions}

以下是从相似图片知识库中检索到的结果供你参考:
{retrieval_context}

用户问题: {user_query}

请用中文回答用户的问题，回答要简洁专业。在回答中，请详细描述你观察到的产品特征，并解释你的判断依据。"""

        # 构建API请求
        payload = {
            "model": self.model_name,
            "prompt": prompt,
            "stream": False,
            "images": [image_b64]
        }
        
        try:
            response = requests.post(f"{self.base_url}/api/generate", json=payload)
            if response.status_code == 200:
                result = response.json()
                return {
                    "success": True,
                    "response": result["response"]
                }
            else:
                return {"success": False, "error": f"API调用失败: HTTP {response.status_code}"}
        except Exception as e:
            return {"success": False, "error": f"调用模型时发生错误: {str(e)}"}
    
    def _format_retrieval_results(self, results: List[Dict]) -> str:
        """
        格式化检索结果
        
        Args:
            results: 检索结果列表
            
        Returns:
            格式化后的检索结果字符串
        """
        if not results:
            return "没有找到相似图片。"
        
        formatted_results = []
        for i, res in enumerate(results):
            label = res.get("label", "未知")
            similarity = res.get("similarity", 0.0)
            path = res.get("path", "")
            
            # 获取文件名
            filename = os.path.basename(path)
            
            formatted_results.append(f"相似图片{i+1}: {filename}, 标签: {label}, 相似度: {similarity:.4f}")
        
        return "\n".join(formatted_results) 
    
    def text_chat(self, user_query: str, product_name: str, knowledge_results: List[Dict]) -> Dict:
        """
        文本问答模式
        
        Args:
            user_query: 用户问题
            product_name: 产品名称
            knowledge_results: 知识库检索结果
            
        Returns:
            对话响应
        """
        # 构建知识库检索结果的描述
        retrieval_context = self._format_retrieval_results(knowledge_results)
        
        # 获取缺陷类型描述
        defect_descriptions = ""
        if hasattr(knowledge_results, "get_defect_features_description"):
            defect_descriptions = knowledge_results.get_defect_features_description(product_name)
        else:
            # 尝试从检索结果中提取缺陷类型
            defect_types = set()
            for result in knowledge_results:
                label = result.get("label", "").strip()
                if label and label != "unknown":
                    defect_types.add(label)
        
            if defect_types:
                defect_descriptions = f"{product_name}产品可能的缺陷类型包括: {', '.join(defect_types)}。"
        
        # 为NEU-DET数据集添加特殊处理
        special_instructions = ""
        if product_name.lower() == "neu-det":
            special_instructions = """
关于钢材表面缺陷的补充知识：

1. 龟裂(crazing)与划痕(scratches)的区别：
   - 龟裂通常是材料内部应力导致的网状裂纹，覆盖一定区域
   - 划痕是外力造成的线状损伤，有明确的方向性

2. 夹杂物(inclusion)与点蚀(pitted_surface)的区别：
   - 夹杂物是冶炼过程中混入的异物，可能呈现为凸起或异色点
   - 点蚀是腐蚀或制造缺陷导致的凹坑，表面有明显的深度变化

3. 斑块(patches)的特征：
   - 表面颜色不均，可能是材料成分不均匀或局部氧化导致
   - 通常没有明显的表面形态变化，主要是颜色或纹理的差异

这些缺陷对钢材性能有不同影响，准确识别对于质量控制至关重要。"""
        
        # 构建提示词
        prompt = f"""你是一个产品表面缺陷检测专家，请回答用户关于{product_name}产品的问题。

{defect_descriptions}

{special_instructions}

以下是从知识库中检索到的相关图片结果供你参考:
{retrieval_context}

用户问题: {user_query}

请用中文回答用户的问题，回答要简洁专业。如果知识库中有相关的缺陷图片，请在回答中详细描述这些缺陷的特征和识别方法。"""

        # 构建API请求
        payload = {
            "model": self.model_name,
            "prompt": prompt,
            "stream": False
        }
        
        try:
            response = requests.post(f"{self.base_url}/api/generate", json=payload)
            if response.status_code == 200:
                result = response.json()
                return {
                    "success": True,
                    "response": result["response"]
                }
            else:
                return {"success": False, "error": f"API调用失败: HTTP {response.status_code}"}
        except Exception as e:
            return {"success": False, "error": f"调用模型时发生错误: {str(e)}"} 